package com.idanchuang.component.shadow.route.sharding;

import com.alibaba.fastjson.JSON;
import com.baomidou.mybatisplus.extension.spring.MybatisSqlSessionFactoryBean;
import com.idanchuang.component.mybatis.plus.ComponentMybatisPlusPureAutoConfiguration;
import com.idanchuang.component.shadow.route.common.aop.ShadowRouteDataSourceAspect;
import com.idanchuang.component.shadow.route.common.conf.ShadowRouteProperties;
import com.idanchuang.component.shadow.route.common.dynamic.ShadowRoutingDataSource;
import com.idanchuang.component.shadow.route.common.filter.ShadowRouteDataSourceFilter;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.session.SqlSessionFactory;
import org.apache.shardingsphere.core.yaml.swapper.ShardingRuleConfigurationYamlSwapper;
import org.apache.shardingsphere.shardingjdbc.api.ShardingDataSourceFactory;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.datasource.ShardingDataSource;
import org.apache.shardingsphere.shardingjdbc.spring.boot.common.SpringBootPropertiesConfigurationProperties;
import org.apache.shardingsphere.shardingjdbc.spring.boot.sharding.SpringBootShardingRuleConfigurationProperties;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.autoconfigure.AutoConfigureAfter;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Map;

/**
 * Shadow Route Sharding Auto Configuration
 *
 * @author Raistlin
 */
@ConditionalOnProperty(prefix = "idanchuang.component.shadow.route", name = "enabled", havingValue = "true")
@AutoConfigureAfter(value = {ComponentMybatisPlusPureAutoConfiguration.class})
@Configuration
public class ComponentShadowRouteShardingAutoConfiguration implements ApplicationContextAware {

    private final Logger logger = LoggerFactory.getLogger(ComponentShadowRouteShardingAutoConfiguration.class.getName());

    private ApplicationContext applicationContext;

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    @Bean
    @ConfigurationProperties(prefix = "idanchuang.component.shadow.route")
    public ShadowRouteProperties shadowRouteProperties() {
        return new ShadowRouteProperties();
    }

    @Bean
    public ShadowRouteDataSourceFilter shadowRouteDataSourceFilter(ShadowRouteProperties shadowRouteProperties) {
        return new ShadowRouteDataSourceFilter(shadowRouteProperties);
    }

    @Bean
    @ConditionalOnMissingBean(ShadowRouteDataSourceFilter.class)
    public ShadowRouteDataSourceAspect shadowRouteDataSourceAspect(ShadowRouteProperties shadowRouteProperties) {
        return new ShadowRouteDataSourceAspect(shadowRouteProperties);
    }

    @Bean
    @Primary
    public DataSource dynamicDataSource(ShadowRouteProperties properties) {
        ShadowRoutingDataSource dynamicDataSource = new ShadowRoutingDataSource();
        Map<Object, Object> dataSources = new HashMap<>();
        if (properties.isUseSharding()) {
            dataSources.put(properties.getActualDataSourceName(), createShardingDataSource(properties.getActualDataSourceName(), properties.getActualDataNodes()));
            dataSources.put(properties.getShadowDataSourceName(), createShardingDataSource(properties.getShadowDataSourceName(), properties.getShadowDataNodes()));
        } else {
            throw new IllegalArgumentException("please config idanchuang.component.shadow.route.use-sharding=true");
        }
        dynamicDataSource.setTargetDataSources(dataSources);
        return dynamicDataSource;
    }

    private DataSource createShardingDataSource(final String dataSourceName, final String dataNodes) {
        DataSource dataSource = applicationContext.getBean(ShardingDataSource.class).getDataSourceMap().get(dataSourceName);
        SpringBootPropertiesConfigurationProperties props = applicationContext.getBean(SpringBootPropertiesConfigurationProperties.class);
        SpringBootShardingRuleConfigurationProperties shardingRule = applicationContext.getBean(SpringBootShardingRuleConfigurationProperties.class);
        SpringBootShardingRuleConfigurationProperties newShardingRule = createShardingRule(shardingRule, dataNodes);
        try {
            return ShardingDataSourceFactory.createDataSource(new LinkedHashMap<String, DataSource>() {{
                put(dataSourceName, dataSource);
            }}, new ShardingRuleConfigurationYamlSwapper().swap(newShardingRule), props.getProps());
        } catch (SQLException e) {
            logger.error("create sharding data source failed", e);
        }
        return null;
    }

    private SpringBootShardingRuleConfigurationProperties createShardingRule(final SpringBootShardingRuleConfigurationProperties shardingRule, final String dataNodes) {
        SpringBootShardingRuleConfigurationProperties result = JSON.parseObject(JSON.toJSONString(shardingRule), SpringBootShardingRuleConfigurationProperties.class);
        result.getTables().forEach((k, v) -> v.setActualDataNodes(dataNodes));
        return result;
    }

    @Bean
    public SqlSessionFactory sqlSessionFactory(@Qualifier("dynamicDataSource") DataSource dynamicDataSource, Interceptor... plugins) throws Exception {
        MybatisSqlSessionFactoryBean sqlSessionFactoryBean = new MybatisSqlSessionFactoryBean();
        sqlSessionFactoryBean.setDataSource(dynamicDataSource);
        sqlSessionFactoryBean.setPlugins(plugins);
        return sqlSessionFactoryBean.getObject();
    }

}
